# stereo_worker.py

from multiprocessing import Process
import numpy as np
import cv2
import os
import time
import traceback
cv2.setNumThreads(4)
from stereo.utils.Filter.fast_least_squares_interpolation import fill_disparity_fast, fill_disparity_morphology

window_size = 5
min_disp = 0
num_disp = 32 - min_disp  # 必须是16的整数倍
def create_sgbm():
    stereo = cv2.StereoSGBM_create(
        minDisparity=min_disp,
        numDisparities=num_disp,
        blockSize=window_size,
        P1=8 * 3 * window_size**2,  # 视差平滑参数
        P2=32 * 3 * window_size**2,
        disp12MaxDiff=1,
        uniquenessRatio=10,
        speckleWindowSize=100,
        speckleRange=64,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY# STEREO_SGBM_MODE_SGBM_3WAY ,STEREO_SGBM_MODE_HH, STEREO_SGBM_MODE_SGBM, STEREO_SGBM_MODE_HH4,STEREO_SGBM_MODE_HH4的速度最快，STEREO_SGBM_MODE_HH的精度最好
    )
    return stereo
sgbm = create_sgbm()
def upscale_disparity(disp_small, scale_x=2, scale_y=2):
    """
    将低分辨率视差图放大，并调整视差值大小
    :param disp_small: 低分辨率视差图（如 320x240）
    :param scale_x: 宽度缩放倍数（如从 320 -> 640 就是 2）
    :param scale_y: 高度缩放倍数
    :return: 放大后的视差图（如 640x480）
    """
    # 1. 插值上采样
    disp_large = cv2.resize(disp_small, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)

    # 2. 缩放视差值（因为位移也是成倍的）
    disp_large *= scale_x  # 通常仅按 x 缩放即可，因为视差是水平位移

    return disp_large

def use_disparity_fast(disp, mode='fast'):
    """
    使用快速最小二乘插值填充视差图空洞

    Args:
        disp: 输入视差图
        mode: 预设模式 ('ultra_fast', 'fast', 'balanced')
        **kwargs: 额外参数

    Returns:
        填充后的视差图
    """
    try:
        filled_disp = fill_disparity_fast(disp, mode=mode)
        # print(f"快速最小二乘插值填充完成，模式: {mode}")
        return filled_disp
    except Exception as e:
        print(f"快速最小二乘插值填充失败，回退到形态学填充: {e}")
        return fill_disparity_morphology(disp)

def estimate_depth(left_image, right_image, down_process=True, copyMake=True):
    """
    进行深度估计推理，返回视差图（深度图）。
    """
    if copyMake:
        left_image = cv2.copyMakeBorder(left_image, 0, 0, num_disp, 0, cv2.BORDER_REPLICATE)
        right_image = cv2.copyMakeBorder(right_image, 0, 0, num_disp, 0, cv2.BORDER_REPLICATE)
    # 转换为灰度图
    if down_process:
        left_image = cv2.resize(left_image, (left_image.shape[1] // 2, left_image.shape[0] // 2), interpolation=cv2.INTER_AREA)
        right_image = cv2.resize(right_image, (right_image.shape[1] // 2, right_image.shape[0] // 2), interpolation=cv2.INTER_AREA)

    # gray_left = cv2.cvtColor(left_image, cv2.COLOR_BGR2GRAY)
    # gray_right = cv2.cvtColor(right_image, cv2.COLOR_BGR2GRAY)

    disp = sgbm.compute(left_image, right_image).astype(np.float32) / 16.0  # SGBM返回的视差需要除以16
    if down_process:
        disp = upscale_disparity(disp, scale_x=2, scale_y=2)  # → 640x480
    if copyMake:
        disp = disp[:, num_disp:]
    disp = use_disparity_fast(disp, mode='fast')
    # disp = cv2.morphologyEx(disp, cv2.MORPH_CLOSE, np.ones((5,5),np.uint8))  # 闭运算填充空洞
    return disp

class StereoMatcherWorker(Process):
    def __init__(self, task_queue, result_queue, shape):
        super().__init__()
        self.task_queue = task_queue
        self.result_queue = result_queue
        self.shape = shape
    
    def run(self):
                # 设置当前子进程的 CPU 亲和性（绑定到大核）
        try:
            core_ids={4, 5, 6, 7}
            pid = os.getpid()
            os.sched_setaffinity(pid, core_ids)
            print(f"[StereoMatcherWorker] PID {pid} 绑定到核心: {core_ids}")
        except Exception as e:
            traceback.print_exc()
            print(f"[StereoMatcherWorker] 设置CPU亲和性失败: {e}")
        while True:
            task = self.task_queue.get()
            if task is None:
                time.sleep(0.1)
                continue
            gray_left_img, gray_right_img, rectified_left, box, job_id = task[0], task[1], task[2], task[3], task[4]
            # disparity = self.stereo.compute(gray_left_img, gray_right_img).astype(np.float32) / 16.0  # SGBM返回的视差需要除以16.0
            disparity = estimate_depth(gray_left_img, gray_right_img, down_process=False, copyMake=True)
            self.result_queue.put((disparity, rectified_left, box, job_id))